package com.uzhie.spring;

import com.alibaba.fastjson.JSONObject;
import org.apache.log4j.Logger;

import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import java.io.*;

/**
 * Created by TangFei on 2018/8/15.
 * httpResponse 包装类
 */
public class HttpServletResponseWrapper extends javax.servlet.http.HttpServletResponseWrapper {

    private Logger logger = Logger.getLogger(getClass());
    private JSONObject responseJSON;

    private ByteArrayOutputStream byteArrayOutputStream;
    private PrintWriter writer;

    {
        this.setCharacterEncoding("utf8");
        responseJSON = new JSONObject();
        byteArrayOutputStream = new ByteArrayOutputStream();
    }

    public HttpServletResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
    }

    public HttpServletResponseWrapper(HttpServletResponse response, HttpServletRequestWrapper httpServletRequestWrapper) throws IOException {
        super(response);
        this.responseJSON.put("requestId", httpServletRequestWrapper.getRequestJSON().getString("requestId"));
        this.responseJSON.put("requestMethod", httpServletRequestWrapper.getRequestJSON().getString("requestMethod"));
    }


    public JSONObject getResponseJSON() {
        logger.info("responseJSON is " + JSONObject.parse(getBytes()));
        return responseJSON;
    }

    @Override
    public javax.servlet.ServletOutputStream getOutputStream() throws IOException {
        logger.info("enter method getOutputStream");
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        return new ServletOutputStream(baos);
    }

    private byte[] getBytes() {
        if(null != writer) {
            writer.close();
            return byteArrayOutputStream.toByteArray();
        }

        if(null != byteArrayOutputStream) {
            try {
                byteArrayOutputStream.flush();
            } catch(IOException e) {
                e.printStackTrace();
            }
        }
        return byteArrayOutputStream.toByteArray();
    }

    private class ServletOutputStream extends javax.servlet.ServletOutputStream {

        private OutputStream outputStream;

        public ServletOutputStream(ByteArrayOutputStream outputStream) {
            this.outputStream = outputStream;
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }

        @Override
        public void write(int b) throws IOException {
            outputStream.write(b);
        }
    }

}
